import unittest
import os
import getpass
import warnings
import ConfigParser
import ast
import csv
import filecmp
import pyfmdb
import re
import psycopg2
import code, traceback

def create_tests(host, port, admin, password, test_update, source_dir, test_dir):

    class FMDBTestCase(unittest.TestCase):
        def setUp(self):
            """
                sets database related attributes taken from the containing factory function returns True or False 
                based on whether the test should test the latest database update script
            """
            self.test_database_update = test_update
            self.host = host
            self.port = port
            self.admin = admin
            self.password = password
            self.source_dir = source_dir
            self.test_dir = test_dir

            if test_update:
                # load repo info from the source directory
                repo = git.Repo(self.source_dir)
                # extract the list of version tags from the repo. assumes all version tags start with 'v'.
                version_tags = [tag for tag in repo.tags if tag.name.startswith('v')]
                # if there are no version tags, rasise an exception
                if not version_tags:
                    raise Exception("Can't test updates. No version tags in git repository.")
                # version tags are expected in the format vx.x.x, where x are integers and there is at least one x.
                # sort the tags by their numbers
                version_tags = sorted(version_tags, key=lambda vtag: tuple([int(n) for n in vtag.name[1:].split('.')]))
                # get the latest tag
                last_tag = version_tags_sorted[-1]
                old_version = {}
                try:
                    # read the contents of VERSION.py from the version associated with the last tag
                    old_versions_def = repo.git.execute('git show ' + last_tag.name + ':VERSION.py')
                except git.GitCommandError:
                    raise Exception("Database update test requested, but git could not retrieve prior version information.")
                else:
                    # extract each version number into old_version
                    for component_version in old_versions_def.split('\n'):
                        component, version = component_version.split('=')
                        old_version[component.strip()] = ast.literal_eval(version)
                    # check that the version of the database in the last version is not the same as the current 
                    # version of the database.
                    if VERSION.db == old_versions['db']:
                        exc = "Database update test requested, but the database version number has not " + \
                              "changed from the last version."
                        raise Exception(exc)

                try:
                    # get the database setup file from the prior version
                    old_db_setup = repo.git.execute('git show ' + tag.name + ':database_setup.sql')
                except git.GitCommandError:
                    raise Exception("Database update test requested, but could not retrieve prior database setup file.")
                else:
                    # create a temporary file and write the old database setup file
                    temp_path = tempfile.mktemp()
                    with open(temp_path, 'w') as f:
                        f.write(old_db_setup.decode('string-escape'))
                    # build the old database
                    pyfmdb.build_database(script=temp_path, host=self.host, port=self.port, admin=self.admin, 
                        password=self.password, delete_existing=True, make_test_user=True)
                    # execute the latest update script
                    update_dir = os.path.join(self.source_dir, 'updates')
                    updates = [name for name in os.listdir(update_dir) if name.startswith('v')]
                    updates = sorted(updates, key=lambda name: [int(x) for x in name.strip('.vsql').split('.')])
                    update = updates[-1]
                    update_path = os.path.join(update_dir, update)

                    pyfmdb.update_database(host=self.host, port=self.port, admin=admin, password=password)
            else:
                pyfmdb.build_database(script=os.path.join(source_dir, 'database_setup.sql'), host=self.host, port=self.port,
                    admin=admin, password=password, delete_existing=True, make_test_user=True)

        def inputdir(self):
            if hasattr(self, 'testdir'):
                return os.path.join(self.testdir, 'input')

        def expectdir(self):
            if hasattr(self, 'testdir'):
                return os.path.join(self.testdir, 'expectations')

    test_cases = []

    @test_cases.append
    class TestCase01(FMDBTestCase):
        name = 'test_import'
        def runTest(self):
            self.testdir = os.path.join(self.test_dir, self.name)

            stage_dirs = ['FlowMonitoring' + str(i) for i in range(6)]
            with pyfmdb.connect(self.host, port=self.port, user='fmdb_tester', password='fmdb_tester') as cnxn:
                for stagedir in stage_dirs:
                    cnxn.batch_update(os.path.join(self.inputdir(), stagedir))
                    outdir = os.path.join(self.expectdir(), stagedir)
                    allflows_exported_path = os.path.join(outdir, 'allflows_exported.csv')
                    with open(allflows_exported_path, 'w') as f:
                        batch = 5000
                        cnxn.execute("SELECT * FROM export_allflows(1)")
                        writer = csv.DictWriter(f, fieldnames=cnxn.cursor_columns, lineterminator='\n')
                        writer.writeheader()
                        while True:
                            rows = cnxn.fetchmany(size=batch, dicts=True)
                            if rows:
                                writer.writerows(rows)
                            else:
                                break

                    allflows_expected_path = os.path.join(outdir, 'allflows_expected.csv')
                    self.assertTrue(filecmp.cmp(allflows_expected_path, allflows_exported_path))

                    eventstats_exported_path = os.path.join(outdir, 'eventstats_exported.csv')
                    with open(eventstats_exported_path, 'w') as f:
                        cnxn.execute("SELECT * FROM export_eventstats(1)")
                        writer = csv.DictWriter(f, fieldnames=cnxn.cursor_columns, lineterminator='\n')
                        writer.writeheader()
                        writer.writerows(cnxn.fetchall(dicts=True))

                    eventstats_expected_path = os.path.join(outdir, 'eventstats_expected.csv')
                    self.assertTrue(filecmp.cmp(eventstats_expected_path, eventstats_exported_path))

                    filesum_exported_path = os.path.join(outdir, 'filesum_exported.csv')
                    with open(filesum_exported_path, 'w') as f:
                        cnxn.execute("SELECT * FROM get_filesummary()")
                        writer = csv.DictWriter(f, fieldnames=cnxn.cursor_columns, lineterminator='\n')
                        writer.writeheader()
                        writer.writerows(cnxn.fetchall(dicts=True))

                    filesum_expected_path = os.path.join(outdir, 'filesum_expected.csv')
                    self.assertTrue(filecmp.cmp(filesum_exported_path, filesum_expected_path))

                    sites_exported_path = os.path.join(outdir, 'sites_exported.csv')
                    with open(sites_exported_path, 'w') as f:
                        cnxn.execute("SELECT * FROM get_sites()")
                        writer = csv.DictWriter(f, fieldnames=cnxn.cursor_columns, lineterminator='\n')
                        writer.writeheader()
                        writer.writerows(cnxn.fetchall(dicts=True))

                    sites_expected_path = os.path.join(outdir, 'sites_expected.csv')

    @test_cases.append
    class TestCase02(FMDBTestCase):
        name = 'test_file_summary'
        def runTest(self):
            self.testdir = os.path.join(self.test_dir, self.name)

            exported_path = os.path.join(self.expectdir(), 'file_summary_exported.csv')
            with pyfmdb.connect(self.host, port=self.port, user='fmdb_tester', password='fmdb_tester') as cnxn:
                cnxn.batch_update(self.inputdir())

                with open(exported_path, 'w') as out:
                    cnxn.execute("SELECT * FROM get_filesummary()")
                    writer = csv.DictWriter(out, fieldnames=cnxn.cursor_columns, lineterminator='\n')
                    writer.writeheader()
                    writer.writerows(cnxn.fetchall(dicts=True))

            expected_path = os.path.join(self.expectdir(), 'file_summary_expected.csv')
            self.assertTrue(filecmp.cmp(exported_path, expected_path))

    @test_cases.append
    class TestCase03(FMDBTestCase):
        name = 'test_excel_formatted'
        def runTest(self):
            self.testdir = os.path.join(self.test_dir, self.name)

            with pyfmdb.connect(self.host, port=self.port, user='fmdb_tester', password='fmdb_tester') as cnxn:
                cnxn.batch_update(os.path.join(self.inputdir(), 'FlowMonitoring'))
                allflows_exported_path = os.path.join(self.expectdir(), 'allflows_exported.csv')
                with open(allflows_exported_path, 'w') as f:
                    batch = 5000
                    cnxn.execute("SELECT * FROM export_allflows(1)")
                    writer = csv.DictWriter(f, fieldnames=cnxn.cursor_columns, lineterminator='\n')
                    writer.writeheader()
                    while True:
                        rows = cnxn.fetchmany(size=batch, dicts=True)
                        if rows:
                            writer.writerows(rows)
                        else:
                            break

                allflows_expected_path = os.path.join(self.expectdir(), 'allflows_expected.csv')
                self.assertTrue(filecmp.cmp(allflows_expected_path, allflows_exported_path))

                eventstats_exported_path = os.path.join(self.expectdir(), 'eventstats_exported.csv')
                with open(eventstats_exported_path, 'w') as f:
                    cnxn.execute("SELECT * FROM export_eventstats(1)")
                    writer = csv.DictWriter(f, fieldnames=cnxn.cursor_columns, lineterminator='\n')
                    writer.writeheader()
                    writer.writerows(cnxn.fetchall(dicts=True))

                eventstats_expected_path = os.path.join(self.expectdir(), 'eventstats_expected.csv')
                self.assertTrue(filecmp.cmp(eventstats_expected_path, eventstats_exported_path))

    @test_cases.append
    class TestCase04(FMDBTestCase):
        name = 'test_user_privileges'
        def runTest(self):
            self.testdir = os.path.join(self.test_dir, self.name)

            outsider_login = 'outsider'
            with pyfmdb.connect(self.host, port=self.port, user=self.admin, password=self.password) as cnxn:
                try:
                    cnxn.execute("DROP USER " + outsider_login)
                except:
                    pass
                cnxn.execute("CREATE ROLE " + outsider_login + " WITH LOGIN PASSWORD '" + outsider_login + "'")

                try:
                    cnxn2 = pyfmdb.connect(self.host, port=self.port, user=outsider_login, password=outsider_login)
                except:
                    outsider_cant_connect = True
                else:
                    outsider_cant_connect = False
                    cnxn2.close()
                self.assertTrue(outsider_cant_connect)


            with pyfmdb.connect(self.host, port=self.port, user='fmdb_tester', password='fmdb_tester') as cnxn:
                cnxn.batch_update(os.path.join(self.inputdir(), 'FlowMonitoring'))
                #print 'bacon'
                #code.interact(local=locals())
                #cnxn.execute("REVOKE CONNECT ON DATABASE fmdb FROM fmdb_user")
                #cnxn.execute("REVOKE CONNECT ON DATABASE fmdb FROM PUBLIC")

            with pyfmdb.connect(self.host, port=self.port, user='fmdb_guest', password='fmdb_guest') as cnxn:
                try:
                    cnxn.execute("SELECT * FROM allflows")
                    cnxn.execute("SELECT * FROM eventstatistics")
                    cnxn.execute("SELECT * FROM filesummary")
                except:
                    traceback.print_exc()
                    can_select = False
                else:
                    can_select = True
                self.assertTrue(can_select)

                try:
                    cnxn.execute("SELECT * FROM private.sites_type")
                except:
                    cant_select_meta = True
                else:
                    cant_select_meta = False
                self.assertTrue(cant_select_meta)

                try:
                    cnxn.execute("SELECT * FROM export_eventstats(1)")
                except:
                    can_execute = False
                else:
                    can_execute = True
                self.assertTrue(can_execute)

    return test_cases

def main():
    source_dir = os.path.dirname(__file__) if '__file__' in globals() else os.path.abspath('.')
    test_dir = os.path.join(source_dir, 'tests')
    cfg_path = os.path.join(test_dir, 'testsuite.ini')

    if os.path.exists(cfg_path):
        parser = ConfigParser.SafeConfigParser()
        parser.read(cfg_path)
        host = ast.literal_eval(parser.get('main', 'host'))
        admin = ast.literal_eval(parser.get('main', 'admin'))
        port = ast.literal_eval(parser.get('main', 'port'))
        test_update = ast.literal_eval(parser.get('main', 'test_update'))
    else:
        host = raw_input('Testing host: ')
        admin = raw_input('Admin login: ')
        while True:
            port = raw_input('Port Number:')
            try:
                port = int(port)
            except:
                print("Invalid input. Port should be an integer.")
            else:
                break
        while True:
            test_update = raw_input('Test latest database update [y/n] >> ').lower()
            if test_update.startswith('y'):
                test_update = True
            elif test_update.startswith('n'):
                test_update = False
            else:
                continue
            break

    password = getpass.getpass('Password for ' + admin + ' on ' + host + ':' + str(port) + ': ')
    tests = create_tests(host=host, port=port, admin=admin, password=password, 
        test_update=test_update, source_dir=source_dir, test_dir=test_dir)

    suite = unittest.TestSuite()
    for test in tests:
        suite.addTest(test())

    return unittest.TextTestRunner(verbosity=2).run(suite)
